/* Copyright (c) 2021, National University of Defense Technology. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "bsr_common.h"
#include "parbsr_laplacian_generator.h"

namespace YHAMG
{

ParBSR_Laplacian_Generator::ParBSR_Laplacian_Generator()
	: nx(100),
	ny(100),
	nz(100),
	bx(2),
	by(2),
	bz(2),
	Px(1),
	Py(1),
	Pz(1),
	cx(1.0),
	cy(1.0),
	cz(1.0)
{
}

ParBSR_Laplacian_Generator::ParBSR_Laplacian_Generator(MPI_Comm comm)
	: ParBSR_Generator(comm),
	nx(100),
	ny(100),
	nz(100),
	bx(2),
	by(2),
	bz(2),
	Px(1),
	Py(1),
	Pz(1),
	cx(1.0),
	cy(1.0),
	cz(1.0)
{
}

ParBSR_Laplacian_Generator::ParBSR_Laplacian_Generator(MPI_Comm comm, int _nx, int _ny, int _nz, int _bx, int _by, int _bz, int _Px, int _Py, int _Pz, double _cx, double _cy, double _cz)
	: ParBSR_Generator(comm),
	nx(_nx),
	ny(_ny),
	nz(_nz),
	bx(_bx),
	by(_by),
	bz(_bz),
	Px(_Px),
	Py(_Py),
	Pz(_Pz),
	cx(_cx),
	cy(_cy),
	cz(_cz)
{
}

void ParBSR_Laplacian_Generator::operator()(ParBSR_Matrix& A) const
{
	int rank;
	MPI_Comm_rank(comm, &rank);

	double d = 0.0;
	if (nx * bx * Px > 1) d -= 2 * cx;
	if (ny * bx * Py > 1) d -= 2 * cy;
	if (nz * bx * Pz > 1) d -= 2 * cz;

	int size = nx * ny * nz;
	int block_size = bx * by * bz;
	int block_nnz = block_size * block_size;

	int* LocLoc_rowptr = new int[size + 1];

	LocLoc_rowptr[0] = 0;
	for (int ix = 0; ix < nx; ++ix)
	{
		for (int iy = 0; iy < ny; ++iy)
		{
			for (int iz = 0; iz < nz; ++iz)
			{
				int cnt = 1;
				if (ix > 0) ++cnt;
				if (iy > 0) ++cnt;
				if (iz > 0) ++cnt;
				if (iz < nz - 1) ++cnt;
				if (iy < ny - 1) ++cnt;
				if (ix < nx - 1) ++cnt;
				LocLoc_rowptr[ix * ny * nz + iy * nz + iz + 1] = cnt;
			}
		}
	}

	for (int i = 0; i < size; ++i)
		LocLoc_rowptr[i + 1] += LocLoc_rowptr[i];

	int* LocLoc_colind = new int[LocLoc_rowptr[size]];
	double* LocLoc_values = new double[LocLoc_rowptr[size] * block_nnz];

	for (int ix = 0; ix < nx; ++ix)
	{
		for (int iy = 0; iy < ny; ++iy)
		{
			for (int iz = 0; iz < nz; ++iz)
			{
				int r = ix * ny * nz + iy * nz + iz;
				int j = LocLoc_rowptr[r];

				if (ix > 0)
				{
					LocLoc_colind[j] = r - ny * nz;
					BSRBlockFill(block_size, 0.0, LocLoc_values + j * block_nnz);

					for (int jy = 0; jy < by; ++jy)
						for (int jz = 0; jz < bz; ++jz)
							LocLoc_values[j * block_nnz + (jy * bz + jz) * block_size + (bx - 1) * by * bz + jy * bz + jz] = cx;
					++j;
				}

				if (iy > 0)
				{
					LocLoc_colind[j] = r - nz;
					BSRBlockFill(block_size, 0.0, LocLoc_values + j * block_nnz);

					for (int jx = 0; jx < bx; ++jx)
						for (int jz = 0; jz < bz; ++jz)
							LocLoc_values[j * block_nnz + (jx * by * bz + jz) * block_size + jx * by * bz + (by - 1) * bz + jz] = cy;
					++j;
				}

				if (iz > 0)
				{
					LocLoc_colind[j] = r - 1;
					BSRBlockFill(block_size, 0.0, LocLoc_values + j * block_nnz);

					for (int jx = 0; jx < bx; ++jx)
						for (int jy = 0; jy < by; ++jy)
							LocLoc_values[j * block_nnz + (jx * by * bz + jy * bz) * block_size + jx * by * bz + jy * bz + bz - 1] = cz;
					++j;
				}

				LocLoc_colind[j] = r;
				BSRBlockFill(block_size, 0.0, LocLoc_values + j * block_nnz);
				
				for (int jx = 0; jx < bx; ++jx)
				{
					for (int jy = 0; jy < by; ++jy)
					{
						for (int jz = 0; jz < bz; ++jz)
						{
							if (jx > 0) LocLoc_values[j * block_nnz + ((jx - 1) * by * bz + jy * bz + jz) * block_size + jx * by * bz + jy * bz + jz] = cx;
							if (jy > 0) LocLoc_values[j * block_nnz + (jx * by * bz + (jy - 1) * bz + jz) * block_size + jx * by * bz + jy * bz + jz] = cy;
							if (jz > 0) LocLoc_values[j * block_nnz + (jx * by * bz + jy * bz + jz - 1) * block_size + jx * by * bz + jy * bz + jz] = cz;
							LocLoc_values[j * block_nnz + (jx * by * bz + jy * bz + jz) * block_size + jx * by * bz + jy * bz + jz] = d;
							if (jz < bz - 1) LocLoc_values[j * block_nnz + (jx * by * bz + jy * bz + jz + 1) * block_size + jx * by * bz + jy * bz + jz] = cz;
							if (jy < by - 1) LocLoc_values[j * block_nnz + (jx * by * bz + (jy + 1) * bz + jz) * block_size + jx * by * bz + jy * bz + jz] = cy;
							if (jx < bx - 1) LocLoc_values[j * block_nnz + ((jx + 1) * by * bz + jy * bz + jz) * block_size + jx * by * bz + jy * bz + jz] = cx;
						}
					}
				}
				
				++j;

				if (iz < nz - 1)
				{
					LocLoc_colind[j] = r + 1;
					BSRBlockFill(block_size, 0.0, LocLoc_values + j * block_nnz);

					for (int jx = 0; jx < bx; ++jx)
						for (int jy = 0; jy < by; ++jy)
							LocLoc_values[j * block_nnz + (jx * by * bz + jy * bz + bz - 1) * block_size + jx * by * bz + jy * bz] = cz;
					++j;
				}

				if (iy < ny - 1)
				{
					LocLoc_colind[j] = r + nz;
					BSRBlockFill(block_size, 0.0, LocLoc_values + j * block_nnz);

					for (int jx = 0; jx < bx; ++jx)
						for (int jz = 0; jz < bz; ++jz)
							LocLoc_values[j * block_nnz + (jx * by * bz + (by - 1) * bz + jz) * block_size + jx * by * bz + jz] = cy;
					++j;
				}

				if (ix < nx - 1)
				{
					LocLoc_colind[j] = r + ny * nz;
					BSRBlockFill(block_size, 0.0, LocLoc_values + j * block_nnz);

					for (int jy = 0; jy < by; ++jy)
						for (int jz = 0; jz < bz; ++jz)
							LocLoc_values[j * block_nnz + ((bx - 1) * by * bz + jy * bz + jz) * block_size + jy * bz + jz] = cx;
					++j;
				}
			}
		}
	}

	int px = rank / (Py * Pz);
	int py = rank / Pz - px * Py;
	int pz = rank - px * Py * Pz - py * Pz;

	int* LocExt_rowptr = new int[size + 1];
	for (int i = 0; i <= size; ++i)
		LocExt_rowptr[i] = 0;

	int num_neighbors = 0;
	int* neighbor_ranks = new int[6];
	int* Ext_ptr = new int[7];

	Ext_ptr[0] = 0;

	if (px > 0)
	{
		for (int iy = 0; iy < ny; ++iy)
			for (int iz = 0; iz < nz; ++iz)
				++LocExt_rowptr[iy * nz + iz];
		neighbor_ranks[num_neighbors++] = rank - Py * Pz;
		Ext_ptr[num_neighbors] = ny * nz;
	}

	if (py > 0)
	{
		for (int ix = 0; ix < nx; ++ix)
			for (int iz = 0; iz < nz; ++iz)
				++LocExt_rowptr[ix * ny * nz + iz];
		neighbor_ranks[num_neighbors++] = rank - Pz;
		Ext_ptr[num_neighbors] = nx * nz;
	}

	if (pz > 0)
	{
		for (int ix = 0; ix < nx; ++ix)
			for (int iy = 0; iy < ny; ++iy)
				++LocExt_rowptr[ix * ny * nz + iy * nz];
		neighbor_ranks[num_neighbors++] = rank - 1;
		Ext_ptr[num_neighbors] = nx * ny;
	}

	if (pz < Pz - 1)
	{
		for (int ix = 0; ix < nx; ++ix)
			for (int iy = 0; iy < ny; ++iy)
				++LocExt_rowptr[ix * ny * nz + iy * nz + nz - 1];
		neighbor_ranks[num_neighbors++] = rank + 1;
		Ext_ptr[num_neighbors] = nx * ny;
	}

	if (py < Py - 1)
	{
		for (int ix = 0; ix < nx; ++ix)
			for (int iz = 0; iz < nz; ++iz)
				++LocExt_rowptr[ix * ny * nz + (ny - 1) * nz + iz];
		neighbor_ranks[num_neighbors++] = rank + Pz;
		Ext_ptr[num_neighbors] = nx * nz;
	}

	if (px < Px - 1)
	{
		for (int iy = 0; iy < ny; ++iy)
			for (int iz = 0; iz < nz; ++iz)
				++LocExt_rowptr[(nx - 1) * ny * nz + iy * nz + iz];
		neighbor_ranks[num_neighbors++] = rank + Py * Pz;
		Ext_ptr[num_neighbors] = ny * nz;
	}

	for (int i = 0; i < size; ++i)
		LocExt_rowptr[i + 1] += LocExt_rowptr[i];

	int* LocExt_colind = new int[LocExt_rowptr[size]];
	double* LocExt_values = new double[LocExt_rowptr[size] * block_nnz];

	for (int r = 0; r < num_neighbors; ++r)
		Ext_ptr[r + 1] += Ext_ptr[r];

	int* Ext_map = new int[Ext_ptr[num_neighbors]];

	int neighbor_index = num_neighbors;

	if (px < Px - 1)
	{
		int t = Ext_ptr[neighbor_index--];
		for (int iy = ny - 1; iy >= 0; --iy)
		{
			for (int iz = nz - 1; iz >= 0; --iz)
			{
				Ext_map[--t] = iy * nz + iz;
				int j = --LocExt_rowptr[(nx - 1) * ny * nz + iy * nz + iz];
				LocExt_colind[j] = t;
				BSRBlockFill(block_size, 0.0, LocExt_values + j * block_nnz);

				for (int jy = 0; jy < by; ++jy)
					for (int jz = 0; jz < bz; ++jz)
						LocExt_values[j * block_nnz + ((bx - 1) * by * bz + jy * bz + jz) * block_size + jy * bz + jz] = cx;
			}
		}
	}


	if (py < Py - 1)
	{
		int t = Ext_ptr[neighbor_index--];
		for (int ix = nx - 1; ix >= 0; --ix)
		{
			for (int iz = nz - 1; iz >= 0; --iz)
			{
				Ext_map[--t] = ix * ny * nz + iz;
				int j = --LocExt_rowptr[ix * ny * nz + (ny - 1) * nz + iz];
				LocExt_colind[j] = t;
				BSRBlockFill(block_size, 0.0, LocExt_values + j * block_nnz);

				for (int jx = 0; jx < bx; ++jx)
					for (int jz = 0; jz < bz; ++jz)
						LocExt_values[j * block_nnz + (jx * by * bz + (by - 1) * bz + jz) * block_size + jx * by * bz + jz] = cy;
			}
		}
	}

	if (pz < Pz - 1)
	{
		int t = Ext_ptr[neighbor_index--];
		for (int ix = nx - 1; ix >= 0; --ix)
		{
			for (int iy = ny - 1; iy >= 0; --iy)
			{
				Ext_map[--t] = ix * ny * nz + iy * nz;
				int j = --LocExt_rowptr[ix * ny * nz + iy * nz + nz - 1];
				LocExt_colind[j] = t;
				BSRBlockFill(block_size, 0.0, LocExt_values + j * block_nnz);

				for (int jx = 0; jx < bx; ++jx)
					for (int jy = 0; jy < by; ++jy)
						LocExt_values[j * block_nnz + (jx * by * bz + jy * bz + bz - 1) * block_size + jx * by * bz + jy * bz] = cz;
			}
		}
	}

	if (pz > 0)
	{
		int t = Ext_ptr[neighbor_index--];
		for (int ix = nx - 1; ix >= 0; --ix)
		{
			for (int iy = ny - 1; iy >= 0; --iy)
			{
				Ext_map[--t] = ix * ny * nz + iy * nz + nz - 1;
				int j = --LocExt_rowptr[ix * ny * nz + iy * nz];
				LocExt_colind[j] = t;
				BSRBlockFill(block_size, 0.0, LocExt_values + j * block_nnz);

				for (int jx = 0; jx < bx; ++jx)
					for (int jy = 0; jy < by; ++jy)
						LocExt_values[j * block_nnz + (jx * by * bz + jy * bz) * block_size + jx * by * bz + jy * bz + bz - 1] = cz;
			}
		}
	}

	if (py > 0)
	{
		int t = Ext_ptr[neighbor_index--];
		for (int ix = nx - 1; ix >= 0; --ix)
		{
			for (int iz = nz - 1; iz >= 0; --iz)
			{
				Ext_map[--t] = ix * ny * nz + (ny - 1) * nz + iz;
				int j = --LocExt_rowptr[ix * ny * nz + iz];
				LocExt_colind[j] = t;
				BSRBlockFill(block_size, 0.0, LocExt_values + j * block_nnz);

				for (int jx = 0; jx < bx; ++jx)
					for (int jz = 0; jz < bz; ++jz)
						LocExt_values[j * block_nnz + (jx * by * bz + jz) * block_size + jx * by * bz + (by - 1) * bz + jz] = cy;
			}
		}
	}

	if (px > 0)
	{
		int t = Ext_ptr[neighbor_index--];
		for (int iy = ny - 1; iy >= 0; --iy)
		{
			for (int iz = nz - 1; iz >= 0; --iz)
			{
				Ext_map[--t] = (nx - 1) * ny * nz + iy * nz + iz;
				int j = --LocExt_rowptr[iy * nz + iz];
				LocExt_colind[j] = t;
				BSRBlockFill(block_size, 0.0, LocExt_values + j * block_nnz);

				for (int jy = 0; jy < by; ++jy)
					for (int jz = 0; jz < bz; ++jz)
						LocExt_values[j * block_nnz + (jy * bz + jz) * block_size + (bx - 1) * by * bz + jy * bz + jz] = cx;
			}
		}
	}

	A.Free();
	A.comm = comm;
	A.LocLoc.size[0] = size;
	A.LocLoc.size[1] = size;
	A.LocLoc.block_size = block_size;
	A.LocLoc.rowptr = LocLoc_rowptr;
	A.LocLoc.colind = LocLoc_colind;
	A.LocLoc.values = LocLoc_values;
	A.LocExt.size[0] = size;
	A.LocExt.size[1] = Ext_ptr[num_neighbors];
	A.LocExt.block_size = block_size;
	A.LocExt.rowptr = LocExt_rowptr;
	A.LocExt.colind = LocExt_colind;
	A.LocExt.values = LocExt_values;
	A.num_neighbors = num_neighbors;
	A.neighbor_ranks = neighbor_ranks;
	A.Ext_ptr = Ext_ptr;
	A.Ext_map = Ext_map;
}

}